package com.wxy.demo.config;

import cn.hutool.cache.Cache;
import cn.hutool.cache.CacheUtil;
import com.google.common.util.concurrent.RateLimiter;
import com.wxy.demo.core.ApiResponse;
import com.wxy.demo.util.IpUtils;
import com.wxy.demo.util.ResponseUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

@Slf4j
public class RateLimiterInterceptor implements HandlerInterceptor {

    private final Cache<String, RateLimiter> limitCache = CacheUtil.newTimedCache(TimeUnit.DAYS.toMillis(1));

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String ip = IpUtils.getIp(request);
        RateLimiter limiter;
        // 判断map集合中是否有创建有创建好的令牌桶
        if (!limitCache.containsKey(ip)) {
            // 创建令牌桶
            limiter = RateLimiter.create(30);
            limitCache.put(ip, limiter);
            log.info("<<=================  请求ip：{},创建令牌桶成功!!!", ip);
        }
        limiter = limitCache.get(ip);
        // 获取令牌
        boolean acquire = limiter.tryAcquire(500, TimeUnit.MILLISECONDS);
        if (!acquire) {
            log.error("服务请求频繁，地址来源：{}", ip);
            ResponseUtils.write(response, ApiResponse.fail("503", "服务器繁忙"));
            return false;
        }
        return true;
    }
}
